# 多项式回归(Polynomial Regression)
import matplotlib.pyplot as plt
import numpy as np
from sklearn import linear_model
from sklearn.preprocessing import PolynomialFeatures

# 读取数据集
# 存储数据中的房屋尺寸和房屋成交价格
datasets_X = []
datasets_Y = []
fr = open('prices.txt', 'r')
# 一次读取整个文件
lines = fr.readlines()
# 逐行进行操作，循环遍历所有数据
for line in lines:
    # 去除数据文件中的逗号
    items = line.strip().split(',')
    datasets_X.append(int(items[0]))
    datasets_Y.append(int(items[1]))

length = len(datasets_X)
# 将datasets_X转化为数组， 并变为二维
datasets_X = np.array(datasets_X).reshape([length, 1])
datasets_Y = np.array(datasets_Y)

minX = min(datasets_X)
maxX = max(datasets_X)
# 以数据datasets_X的最大值和最小值为范围，建立等差数列，方便后续画图
X = np.arange(minX, maxX).reshape([-1, 1])

# degree=2表示建立datasets_X的二 次多项式特征X_poly
poly_reg = PolynomialFeatures(degree=2)
# 创建线性回归，使用线性模型学习X_poly和datasets_Y之间的映射关系（即参数）
X_poly = poly_reg.fit_transform(datasets_X)
lin_reg_2 = linear_model.LinearRegression()
lin_reg_2.fit(X_poly, datasets_Y)

# 图像中显示
plt.scatter(datasets_X, datasets_Y, color='red')
# 先将X处理成 多项式特征
plt.plot(X, lin_reg_2.predict(poly_reg.fit_transform(X)), color='blue')
plt.xlabel('Area')
plt.ylabel('Price')
plt.show()
